'''
前面已经划分了3D的nii文件的训练集和测试集，现在要把2D切片的训练集和测试集id保存到txt文件中
'''
import os
import sys
import numpy as np
sys.path.append(os.path.split(sys.path[0])[0])
from config import zunYiParameter as para
from sklearn.model_selection import train_test_split

def process(nii3d_train_id_list,nii3d_test_id_list):
    train_2d_id_list=[]
    test_2d_id_list = []
    #volume/33_G004.nii.gz segmentation/33_G004.nii.gz
    nii_2d_id_list = os.listdir(para.cut2d_save_path_vol)
    for nii_2d_id in nii_2d_id_list:
        is_train_id=True
        name=nii_2d_id.split('_')[0]#结果为例如‘G027’
        for test_2d_id in nii3d_test_id_list:
            if name in test_2d_id:
                is_train_id=False
        name=nii_2d_id.split('.')[0]
        id_name = os.path.join('volume', name + '.nii.gz') \
                  + ' ' + os.path.join('segmentation', name + '.nii.gz')#例如'volume/G014.nii.gz segmentation/G014.nii.gz'
        if is_train_id==True:
            train_2d_id_list.append(id_name)
        else:
            test_2d_id_list.append(id_name)
    print('训练集和测试集的2D切片数量分别为{}，{}'.format(len(train_2d_id_list),len(test_2d_id_list)))
    if not os.path.exists(os.path.dirname(para.train2d_choose_id_path)):
        os.makedirs(os.path.dirname(para.train2d_choose_id_path))
    if not os.path.exists(os.path.dirname(para.test2d_id_path)):
        os.makedirs(os.path.dirname(para.test2d_id_path))
    np.savetxt(para.train2d_choose_id_path, np.array(train_2d_id_list), fmt = "%s")
    np.savetxt(para.test2d_id_path, np.array(test_2d_id_list), fmt = "%s")
if __name__ == '__main__':
    with open(para.train_nii_id_path, 'r') as f:
        train_id_list = f.read().splitlines()
    with open(para.test_nii_id_path, 'r') as f:
        test_id_list = f.read().splitlines()

    process(train_id_list,test_id_list)